import numpy as np
import tkinter as tk
from tkinter import ttk
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

# 设置支持中文的字体并放大
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['font.size'] = 14

# 四元素类
class Quaternion:
    def __init__(self, w, x, y, z):
        self.w = w
        self.x = x
        self.y = y
        self.z = z
    
    def __str__(self):
        return f"{self.w:.2f} + {self.x:.2f}i + {self.y:.2f}j + {self.z:.2f}k"
    
    def multiply(self, other):
        w1, x1, y1, z1 = self.w, self.x, self.y, self.z
        w2, x2, y2, z2 = other.w, other.x, other.y, other.z
        w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
        x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
        y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
        z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
        return Quaternion(w, x, y, z)
    
    def conjugate(self):
        return Quaternion(self.w, -self.x, -self.y, -self.z)
    
    def rotate_vector(self, vector):
        p = Quaternion(0, vector[0], vector[1], vector[2])
        q_conj = self.conjugate()
        result = self.multiply(p).multiply(q_conj)
        return np.array([result.x, result.y, result.z])

# 创建四元素
def create_quaternion(axis, angle_deg):
    angle_rad = np.radians(angle_deg)
    w = np.cos(angle_rad / 2)
    s = np.sin(angle_rad / 2)
    norm = np.sqrt(sum(a**2 for a in axis))
    if norm == 0:
        return Quaternion(1, 0, 0, 0)
    x, y, z = [a / norm * s for a in axis]
    return Quaternion(w, x, y, z)

# 可视化工具类
class QuaternionVisualizer:
    def __init__(self, root):
        self.root = root
        self.root.title("四元素可视化工具")
        self.root.geometry("900x700")
        
        # 设置全局字体
        self.font_large = ('SimHei', 14)
        
        # 初始向量
        self.initial_vector = np.array([1, 0, 0])
        self.rotated_vector = self.initial_vector.copy()
        self.current_q = Quaternion(1, 0, 0, 0)
        
        # 主框架
        self.frame = ttk.Frame(root, padding="20")
        self.frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 设置行列权重
        self.root.grid_rowconfigure(0, weight=1)
        self.root.grid_columnconfigure(0, weight=1)
        self.frame.grid_rowconfigure(0, weight=1)
        self.frame.grid_columnconfigure(2, weight=1)

        # 输入区域（左侧）
        self.input_frame = ttk.Frame(self.frame)
        self.input_frame.grid(row=0, column=0, sticky=(tk.W, tk.N, tk.S), padx=10)
        
        ttk.Label(self.input_frame, text="旋转轴 (x y z):", font=self.font_large).grid(row=0, column=0, sticky=tk.W, pady=10)
        self.axis_entry = ttk.Entry(self.input_frame, width=15, font=self.font_large)
        self.axis_entry.grid(row=0, column=1, sticky=(tk.W, tk.E), pady=10)
        self.axis_entry.insert(0, "0 0 1")

        ttk.Label(self.input_frame, text="旋转角度 (度):", font=self.font_large).grid(row=1, column=0, sticky=tk.W, pady=10)
        self.angle_entry = ttk.Entry(self.input_frame, width=15, font=self.font_large)
        self.angle_entry.grid(row=1, column=1, sticky=(tk.W, tk.E), pady=10)
        self.angle_entry.insert(0, "90")

        ttk.Button(self.input_frame, text="应用旋转", command=self.apply_rotation, style='Large.TButton').grid(row=2, column=0, columnspan=2, pady=10)
        ttk.Button(self.input_frame, text="重置", command=self.reset, style='Large.TButton').grid(row=3, column=0, columnspan=2, pady=10)

        # 四元素信息（中部）
        self.info_frame = ttk.Frame(self.frame)
        self.info_frame.grid(row=3, column=0, sticky=(tk.W, tk.N, tk.S), padx=10)
        
        self.q_label = ttk.Label(self.info_frame, text="当前四元素: 1.00 + 0.00i + 0.00j + 0.00k", font=self.font_large)
        self.q_label.grid(row=4, column=0, pady=10)
        
        self.double_cover_label = ttk.Label(self.info_frame, text="双覆盖版本: -1.00 + 0.00i + 0.00j + 0.00k", font=self.font_large)
        self.double_cover_label.grid(row=5, column=0, pady=10)
        
        self.explain_label = ttk.Label(self.info_frame, text="解释: 未旋转状态", font=self.font_large, wraplength=400)
        self.explain_label.grid(row=6, column=0, pady=10)

        # Matplotlib 嵌入（右侧）
        self.fig = plt.Figure(figsize=(7, 7))
        self.ax = self.fig.add_subplot(111, projection='3d')
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.frame)
        self.canvas.get_tk_widget().grid(row=0, column=2, sticky=(tk.W, tk.E, tk.N, tk.S))

        # 绑定窗口大小变化事件
        self.root.bind("<Configure>", self.on_resize)
        
        # 自定义按钮样式
        style = ttk.Style()
        style.configure('Large.TButton', font=self.font_large)
        
        # 初次绘制
        self.update_plot()
        self.last_width = self.root.winfo_width()
        self.last_height = self.root.winfo_height()

    def update_plot(self):
        self.ax.clear()
        self.ax.quiver(0, 0, 0, self.initial_vector[0], self.initial_vector[1], self.initial_vector[2], 
                       color='r', label='初始向量', linewidth=2)
        self.ax.quiver(0, 0, 0, self.rotated_vector[0], self.rotated_vector[1], self.rotated_vector[2], 
                       color='b', label='旋转后向量', linewidth=2)
        self.ax.set_xlim([-1.5, 1.5])
        self.ax.set_ylim([-1.5, 1.5])
        self.ax.set_zlim([-1.5, 1.5])
        self.ax.set_xlabel('X轴')
        self.ax.set_ylabel('Y轴')
        self.ax.set_zlabel('Z轴')
        self.ax.legend()
        self.canvas.draw()

    def on_resize(self, event):
        # 避免初始调整和重复绘制
        if event.widget == self.root:
            new_width = event.width
            new_height = event.height
            if new_width != self.last_width or new_height != self.last_height:
                # 只调整图形大小，不重绘内容
                dpi = self.fig.dpi
                self.fig.set_size_inches((new_width * 0.5) / dpi, (new_height * 0.8) / dpi)
                self.canvas.draw()  # 只更新画布，不重新生成视图
                self.last_width = new_width
                self.last_height = new_height

    def apply_rotation(self):
        try:
            axis = [float(x) for x in self.axis_entry.get().split()]
            if len(axis) != 3:
                raise ValueError("轴必须是3个数字")
            angle = float(self.angle_entry.get())
        except ValueError as e:
            tk.messagebox.showerror("输入错误", f"无效输入: {e}")
            return
        
        q = create_quaternion(axis, angle)
        self.current_q = q.multiply(self.current_q)
        self.rotated_vector = self.current_q.rotate_vector(self.initial_vector)
        
        self.q_label.config(text=f"当前四元素: {self.current_q}")
        self.double_cover_label.config(text=f"双覆盖版本: {Quaternion(-self.current_q.w, -self.current_q.x, -self.current_q.y, -self.current_q.z)}")
        
        norm_axis = np.array(axis) / np.sqrt(sum(a**2 for a in axis))
        explain_text = (
            f"解释:\n"
            f"1. 四元素 = w + xi + yj + zk\n"
            f"   - w = cos(θ/2) = {self.current_q.w:.2f} (标量部分，表示旋转角度)\n"
            f"   - (x, y, z) = ({self.current_q.x:.2f}, {self.current_q.y:.2f}, {self.current_q.z:.2f}) (向量部分，表示轴)\n"
            f"2. 当前旋转轴: [{norm_axis[0]:.2f}, {norm_axis[1]:.2f}, {norm_axis[2]:.2f}]\n"
            f"3. 双覆盖: q 和 -q 表示相同旋转"
        )
        self.explain_label.config(text=explain_text)
        
        self.update_plot()

    def reset(self):
        self.current_q = Quaternion(1, 0, 0, 0)
        self.rotated_vector = self.initial_vector.copy()
        
        self.axis_entry.delete(0, tk.END)
        self.axis_entry.insert(0, "0 0 1")
        self.angle_entry.delete(0, tk.END)
        self.angle_entry.insert(0, "90")
        
        self.q_label.config(text="当前四元素: 1.00 + 0.00i + 0.00j + 0.00k")
        self.double_cover_label.config(text="双覆盖版本: -1.00 + 0.00i + 0.00j + 0.00k")
        self.explain_label.config(text="解释: 未旋转状态")
        
        self.update_plot()

# 启动工具
if __name__ == "__main__":
    root = tk.Tk()
    app = QuaternionVisualizer(root)
    root.mainloop()